This is a really cool PR, I don't think I've seen anyone recently manually manage process pools as opposed to just calling torchrun. I think this solution works quite nicely for us though because we don't care (yet) about multi node or fault tolerance

Idk if he has time but just in case @kiukchung would really appreciate your feedback on whether this is a reasonable way of supporting a job system where people submit distributed kernels

The part I'm a bit worried about is the example right now can check that a kernel is globally correct by iterating over many ranks and making sure they're locally correct. That's probably not going to be true for any of the distributed AMD kernels. This is fine as long as we offload the burden of gathering all the outputs onto the user

my first attempt was to just torchrun eval.py, but that breaks our file descriptor approach for result reporting. so i followed @danielhua23's example and did the process startup manually, which fit quite nicely with the existing code that had a process pool of one process anyway.

If you cannot do the check locally, there's nothing preventing us from having all-gather calls inside the validation function.

Idk if he has time but just in case @kiukchung would really appreciate your feedback on whether this is a reasonable way of supporting a job system where people submit distributed kernels

Managing your own processes is perfectly valid. torchrun handles things like: error propagation, proces-group-gang-semantics, better logging/console output, etc. That said, it was primarily built for training applications.

The use-case here seems to be data-parallel eval in a single host (but multi-gpu) where the rank is used to partition the eval dataset, so perhaps gang-semantics is not desired (e.g. torchrun will SIGTERM all workers if at least 1 worker fails, which is the behavior we want during training).

For completeness I'll mention that torchrun can be used programmatically (see: https://github.com/pytorch/pytorch/blob/main/torch/distributed/run.py#L927-L930):

from torch.distributed.launcher.api import elastic_launch, LaunchConfig, DefaultLogsSpecs, Std

config = LaunchConfig(
   role=f"eval_{name_of_task}", # name the worker procs (mostly informational)
   min_nodes=1,
   max_nodes=1, # set both min and max to 1 in your case (single host)
   nproc_per_node=num_gpus, 
   rdzv_endpoint="127.0.0.1",
   rdzv_backend="c10d",
   max_restarts=0, # do not restart the gang (fail on the first failure)
)
  
elastic_launch(config, entrypoint=main_fn)(main_fn_args)

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
src/libkernelbot
consts.py
task.py					119, 162
Project Total

This report was generated by python-coverage-comment-action